import math

import numpy as np
import cv2 as cv


def mosaic_template_match(search_img, template_img, loc, *, 
    mosaic_width=16
):
    height, width = search_img.shape[0], search_img.shape[1]
    t_height, t_width = template_img.shape[0], template_img.shape[1]

    template = np.zeros_like(search_img)
    template[loc[1]:loc[1]+t_height, loc[0]:loc[0]+t_width] = template_img

    mask_row_num = math.ceil(height/mosaic_width)
    mask_col_num = math.ceil(width/mosaic_width)

    row_mask = []
    flag = True
    for mr in range(mask_row_num):
        for i in range(mosaic_width):
            row_mask.append(flag)
        flag = (not flag)

    col_mask = []
    flag = True
    for mc in range(mask_col_num):
        for i in range(mosaic_width):
            col_mask.append(flag)
        flag = (not flag)

    row_mask = row_mask[:height]
    col_mask = row_mask[:width]
    row_mask = np.expand_dims(np.array(row_mask), axis=1)
    col_mask = np.expand_dims(np.array(col_mask), axis=0)

    mosaic_mask = (row_mask == col_mask)
    mosaic_neg_mask = np.logical_not(mosaic_mask)
    mosaic_img = np.zeros_like(search_img)

    mosaic_img[mosaic_mask] = search_img[mosaic_mask]
    mosaic_img[mosaic_neg_mask] = template[mosaic_neg_mask]

    return mosaic_img


if __name__ == "__main__":

    # template_img = cv.imread("E:/datasets/client-data/match-B/os-val/resddfn3/img/sar=img-10-5.png")
    # search_img = cv.imread("E:/datasets/client-data/match-B/os-val/resddfn3/img/opt=img-10-5.png")
    # mosaic = mosaic_template_match(search_img, template_img, [38, 38])

    # # print(mosaic.shape)
    # cv.imwrite("E:/workspace/SOMatch/tmp/image/mosaic.png", mosaic)

    loc = {
        # "INTEG" : [38, 38], 
        # "sslc" : [40, 38], 
        # "ddfn" : [41, 39], 
        # "unet" : [43, 37], 
        "kw" : [140, 134]
    }

    template_img = cv.imread("C:/Users/FU/Desktop/sar256.png")
    search_img = cv.imread("C:/Users/FU/Desktop/opt512.png")

    for k, v in loc.items():
        mosaic = mosaic_template_match(search_img, template_img, v, mosaic_width=56)
        cv.imwrite(f"C:/Users/FU/Desktop/simulate-mosaic.png", mosaic)

